// YemuMalloc
#include <stddef.h>
#include "YemuEnv.h"

// Structure for a free block in the heap
struct HeapFreeBlock {
    unsigned long size;                // Size of this free block (bytes, excluding header)
    struct HeapFreeBlock* next;        // Pointer to next free block
};

struct HeapFreeBlock* HeapFreeBlockHead = NULL; // Head pointer of the free list

// Initialize the heap: set up the first big free block
void YemuMallocInit(unsigned long heap_start, unsigned long heap_end) {
    YemuPrint("Heap : [0x%lx - 0x%lx)\n", heap_start, heap_end);
    // The addr of HeapFreeBlockHead(which is a point of struct HeapFreeBlock) is the heap_start
    HeapFreeBlockHead = (struct HeapFreeBlock*)heap_start;
    HeapFreeBlockHead->size = heap_end - heap_start - sizeof(struct HeapFreeBlock);
    HeapFreeBlockHead->next = NULL;
    return ;
}

// Allocate memory from the heap
void* YemuMalloc(unsigned int size) {
    // Align to 8 bytes
    size = (size + 7) & ~7;

    struct HeapFreeBlock* prev = NULL;
    struct HeapFreeBlock* curr = HeapFreeBlockHead;

    // Traverse the free list to find a suitable block
    while (curr != NULL) {
        if (curr->size >= size) {
            // If the block is much bigger, split it
            if (curr->size >= size + sizeof(struct HeapFreeBlock) + 8) {
                // New free block after allocation
                unsigned long new_block_addr = (unsigned long)curr + sizeof(struct HeapFreeBlock) + size;
                struct HeapFreeBlock* new_block = (struct HeapFreeBlock*)new_block_addr;
                new_block->size = curr->size - size - sizeof(struct HeapFreeBlock);
                new_block->next = curr->next;

                // Update the current block
                curr->size = size;
                if (prev) prev->next = new_block;
                else HeapFreeBlockHead = new_block;

                // Return pointer to user memory after header
                return (void*)((unsigned long)curr + sizeof(struct HeapFreeBlock));
            } else {
                // Use the whole block (not enough space to split)
                if (prev) prev->next = curr->next;
                else HeapFreeBlockHead = curr->next;
                return (void*)((unsigned long)curr + sizeof(struct HeapFreeBlock));
            }
        }
        prev = curr;
        curr = curr->next;
    }
    // No suitable block found
    return NULL;
}

// Free previously allocated memory and merge blocks if possible
void YemuFree(void* ptr) {
    if (!ptr) return;

    // Get pointer to the block header
    struct HeapFreeBlock* block = (struct HeapFreeBlock*)((unsigned long)ptr - sizeof(struct HeapFreeBlock));
    struct HeapFreeBlock* prev = NULL;
    struct HeapFreeBlock* curr = HeapFreeBlockHead;

    // Insert the block back into the free list in address order
    while (curr != NULL && curr < block) {
        prev = curr;
        curr = curr->next;
    }
    block->next = curr;
    if (prev) prev->next = block;
    else HeapFreeBlockHead = block;

    // Merge with previous block if adjacent
    if (prev != NULL && (unsigned long)prev + sizeof(struct HeapFreeBlock) + prev->size == (unsigned long)block) {
        prev->size += sizeof(struct HeapFreeBlock) + block->size;
        prev->next = block->next;
        // After merging with the previous block, update "block" 
        // to point to the merged block so that we can correctly 
        // check and merge with the next block as well.
        block = prev;
    }
    // Merge with next block if adjacent
    if (curr != NULL && (unsigned long)block + sizeof(struct HeapFreeBlock) + block->size == (unsigned long)curr) {
        block->size += sizeof(struct HeapFreeBlock) + curr->size;
        block->next = curr->next;
    }

    return ;
}

void YemuHeapStatus() {
    struct HeapFreeBlock* curr = HeapFreeBlockHead;
    int block_cnt = 0;
    YemuPrint("\n========== Heap Status ==========\n");
    YemuPrint("Heap Free Block list :\n");
    curr = HeapFreeBlockHead;
    while (curr != NULL) {
        // Output Info : Block NUM and [0xADDR:SIZE]
        YemuPrint("Block %d : [0x%lx:%u]", block_cnt++, (unsigned long)curr, curr->size);
        if (curr->next != NULL) YemuPrint(" -> \n");
        curr = curr->next;
    }
    YemuPrint("\n=================================\n\n");
}

void YemuHeapTest() {
    YemuPrint("==== Complex Stress Test Start ====\n\n");

    YemuPrint("Step 0: Check init status.\n");
    YemuHeapStatus();

    const int N = 20;
    void* ptrs[N];

    YemuPrint("Step 1: Allocate 20 small blocks.\n");
    for (int i = 0; i < N; i++) {
        ptrs[i] = YemuMalloc(64 + (i % 4) * 32); // sizes alternate: 64, 96, 128, 160
        YemuPrint("Alloc[%d] = %lx (size=%d)\n", i, (unsigned long)ptrs[i], 64 + (i % 4) * 32);
    }
    YemuHeapStatus();

    YemuPrint("Step 2: Free blocks with even indices to create fragmentation.\n");
    for (int i = 0; i < N; i += 2) {
        YemuFree(ptrs[i]);
        YemuPrint("Free[%d] = %lx\n", i, (unsigned long)ptrs[i]);
    }
    YemuHeapStatus();

    YemuPrint("Step 3: Allocate smaller blocks into the fragmented space.\n");
    void* newptrs[10];
    for (int i = 0; i < 10; i++) {
        newptrs[i] = YemuMalloc(48 + (i % 3) * 24); // 48,72,96
        YemuPrint("Re-Alloc[%d] = %lx (size=%d)\n", i, (unsigned long)newptrs[i], 48 + (i % 3) * 24);
        YemuHeapStatus();
    }
    
    YemuPrint("Step 4: Free all old blocks with odd indices.\n");
    for (int i = 1; i < N; i += 2) {
        YemuFree(ptrs[i]);
        YemuPrint("Free-old[%d] = %lx\n", i, (unsigned long)ptrs[i]);
        YemuHeapStatus();
    }

    YemuPrint("Step 5: Partially free newly allocated blocks (free the first 5).\n");
    for (int i = 0; i < 5; i++) {
        YemuFree(newptrs[i]);
        YemuPrint("Free-new[%d] = %lx\n", i, (unsigned long)newptrs[i]);
        YemuHeapStatus();
    }

    YemuPrint("Step 6: Allocate larger blocks to test fragmented space utilization.\n");
    void* big1 = YemuMalloc(512);
    YemuPrint("Big1=%lx(512)\n", (unsigned long)big1);
    YemuHeapStatus();
    void* big2 = YemuMalloc(1024);
    YemuPrint("Big2=%lx(1024)\n", (unsigned long)big2);
    YemuHeapStatus();

    YemuPrint("Step 7: Free all remaining new blocks.\n");
    for (int i = 5; i < 10; i++) {
        YemuFree(newptrs[i]);
        YemuPrint("Free-new[%d] = %lx\n", i, (unsigned long)newptrs[i]);
        YemuHeapStatus();
    }

    YemuPrint("Freed big1\n");
    YemuFree(big1);
    YemuHeapStatus();

    YemuPrint("Freed big2\n");
    YemuFree(big2);
    YemuHeapStatus();

    YemuPrint("==== Complex Stress Test End ====\n");
    return ;
}
